# -*- coding: utf-8 -*-


# 制作人脸特征向量的数据库 最后会保存两个文件，分别是: 1 据库中的人脸特征向量 2 对应的名字
# MT-CNN: 人脸检测   ：   faceNet: 人脸识别 （效果很一般，考虑后期换成InsightFace）

# 1.通过MT-CNN人脸检测模型，从照片中提取人脸图像。
# 2.把人脸图像输入到FaceNet，计算Embedding的特征向量。
# 3.比较特征向量间的欧式距离，判断是否为同一人，例如当特征距离小于1的时候认为是同一个人，特征距离大于1的时候认为是不同人。

from facenet_pytorch import MTCNN, InceptionResnetV1
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
import numpy as np
import os

workers = 0 if os.name == 'nt' else 4
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# names 是库中的 person name， embeddings 是对应的512维的人脸特征向量
names_path = "./database/names.pt"
embeddings_path = "./database/embeddings.pt"

# 1.加载并定义要用于人脸检测的 MTCNN 检测网络
# 根据源码中注释
# image_size:  输出图像大小（以px为单位）正方形  会发现输出保存的人脸都是正方形的
mtcnn = MTCNN(
    image_size=160, margin=0, min_face_size=20,
    thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
    device=device
)

# 2. 加载在 vggface2 人脸数据集上预训练好的 InceptionResnetV1 网络;
# InceptionResnetV1 提供了两个预训练模型，分别在 vggface 数据集和 casia 数据集上训练的。
# 然后放到C:\Users\你的用户名\.cache\torch\checkpoints这个文件夹下面
# 如果是linux系统，那么存放在/home/你的用户名/.cache/torch/checkpoints下面
resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)


# InceptionResnetV1 则是用来进行人脸特征提取获得特征编码，进而进行特征相似度计算


# 定义读取数据集，类别名称与索引转换的函数
def collate_fn(x):
    return x[0]


# 返回列表 List 中 target 出现的所有下标
def find_index(needlist, target):
    index = []
    for i in range(len(needlist)):
        if needlist[i] == target:
            index.append(i)
    return index


def main():
    # 将所有的单人照图片放在各自的文件夹中，文件夹名字就是人的名字,存放格式如下
    """
    --orgin
      |--zhangsan
         |--1.jpg
         |--2.jpg
      |--lisi
         |--1.jpg
         |--2.jpg
    """

    dataset = datasets.ImageFolder('./database/orgin')  # 加载人脸特征和名字数据库

    # 此处class为人脸类别的list，可自行打印查看
    dataset.idx_to_class = {i: c for c, i in dataset.class_to_idx.items()}
    loader = DataLoader(dataset, collate_fn=collate_fn, num_workers=workers)
    print(dataset.idx_to_class)
    aligned = []  # aligned 就是从图像上抠出的人脸
    names = []
    i = 1

    for x, y in loader:
        path = './database/aligned/{}/'.format(dataset.idx_to_class[y])  # 这个是要保存的人脸路径
        if not os.path.exists(path):
            print('文件夹新建中...')
            i = 1
            os.makedirs(path)
        # 这里遍历i，因为类别中的名字是按数字正序命名的
        x_aligned, prob, _ = mtcnn(x, return_prob=True, save_path=path + '/{}.jpg'.format(i))
        i = i + 1
        print(_)
        if x_aligned is not None:
            # print('Face detected with probability: {:8f}'.format(prob))
            aligned.append(x_aligned)
            names.append(dataset.idx_to_class[y])
    aligned = torch.stack(aligned).to(device)
    embeddings = resnet(aligned).detach().cpu()  # 提取所有人脸的特征向量，每个向量的长度是512

    # 对 name 中重复的进行处理，重复的取均值
    # names_new -> 出现过的类别
    names_new = []
    embeddings_new = []
    embeddings_list = embeddings.numpy()
    for i in range(len(names)):
        if names[i] not in names_new:
            names_new.append(names[i])
            # index是各个类别中出现过的图片的数组下标
            index = find_index(names, names[i])
            s = len(index)
            # 归一化后的人脸特征向量
            embeddings_new.append(embeddings_list[i] * 0)
            for j in range(s):
                embeddings_new[-1] += embeddings_list[index[j]] / s
        else:
            pass

    embeddings_new = torch.Tensor(embeddings_new)
    embeddings_new.requires_grad_(False)

    torch.save(embeddings_new, embeddings_path)  # 保存在一个文件
    torch.save(names_new, names_path)


if __name__ == '__main__':
    main()
